from termcolor import colored
from ultralytics import YOLO#type: ignore
import cv2
import os
import time
import json
from DepthColorProcessor import DepthColorProcessor
# modelDetectPath = rf"{MODEL_ROOT}\yolov8n.pt"
modelDetectPath = "best13.pt"
# modelDetectPath = "best13_ncnn_model"
print("正在加载模型文件")
modelDetect = YOLO(modelDetectPath)  # 通常是pt模型的文件
print("模型加载成功！")
dcp = DepthColorProcessor(interactive_plots=False)
# 打开摄像头
def main():
    frame = None
    while True:
        current_frame = dcp.color_bin_to_png("color_data.bin", "yolo.png")
        frame = current_frame.copy()
        resultsDetect = modelDetect.predict(
            source="yolo.png",
            imgsz=[640, 480],  # 此处可以调节
            half=False,
            # iou=0.5,
            conf=0.1,
            verbose=False,  # 添加此参数来抑制输出
        )

        # 解析检测结果
        annotated_frame = resultsDetect[0].plot()  # 自动绘制检测结果
        yolo_obj = resultsDetect[0]
        
        # 提取检测到的对象信息（边界框坐标和类别名称）
        detected_objects = []
        for box in yolo_obj.boxes:
            # 获取边界框坐标 (x1, y1, x2, y2 格式)
            coords = box.xyxy[0].tolist()  # 转换为列表
            # 获取置信度
            conf = float(box.conf[0])
            # 获取类别ID和名称
            class_id = int(box.cls[0])
            class_name = yolo_obj.names[class_id]
            detected_objects.append({
                'pos': [(coords[0] + coords[2]) / 2, (coords[1] + coords[3]) / 2],  # [center_x, center_y]
                'label_id': class_id,
                'label': class_name,
                'conf': conf,
            })
        with open("detected_objects.json", "w") as f:
            json.dump({"objs":detected_objects, "time":time.time()}, f, indent=4)
        cv2.imshow('YOLOv8', annotated_frame)
        cv2.waitKey(1)

if __name__ == "__main__":
    main()
    # for yolo_obj in main():
    #     print(yolo_obj)
